from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory, RunnableConfig
from langchain_openai import ChatOpenAI

from env_utils import LOCAL_API_KEY, LOCAL_BASE_URL, DEEPSEEK_API_KEY, DEEPSEEK_BASE_URL
import gradio as gr

prompt = ChatPromptTemplate.from_messages([
    ('system', '你是一个资深动漫专家，根据用户输入以二次元的语气回复相应的输出'),
    MessagesPlaceholder(variable_name="chat_history", optional=True),
    ('human', '{input}')
])

# llm = ChatOpenAI(
#     model="qwen3-8b",
#     temperature=0.8,
#     api_key=LOCAL_API_KEY,
#     base_url=LOCAL_BASE_URL,
#     extra_body={'chat_template_kwargs': {'enable_thinking': True}},
# )
llm = ChatOpenAI(
    model="deepseek-chat",
    temperature=0.8,
    api_key=DEEPSEEK_API_KEY,
    base_url=DEEPSEEK_BASE_URL
)

chain = prompt | llm

store = {}

#存储历史消息/InMemoryChatMessageHistory存在内存,数据库存储SQLChatMessageHistory,Redis存储RedisChatMessageHistory
def get_session_history(session_id: str):
    if session_id not in store:
        store[session_id] = InMemoryChatMessageHistory()
        from langchain_community.chat_message_histories import SQLChatMessageHistory
        from langchain_community.chat_message_histories import RedisChatMessageHistory
    return store[session_id]


chain_with_message_history = RunnableWithMessageHistory(
    chain,
    get_session_history,
    input_messages_key='input',
    history_messages_key='chat_history'
)

#剪辑内容，生成摘要
def summarize_messages(current_input):
    session_id = current_input['config']['configurable']['session_id']
    if not session_id:
        raise ValueError("必须通过config参数提供session_id")
    chat_history = get_session_history(session_id)
    stored_messages = chat_history.messages
    if len(stored_messages) < 2:
        return False
    #剪辑消息
    last_two_messages = stored_messages[-2:]
    messages_to_summarize = stored_messages[:-2]
    summarize_prompt = ChatPromptTemplate.from_messages([
        ('system', '请将一下历史压缩为一条保留关键信息的摘要信息'),
        ('placeHolder', '{chat_history}'),
        ('human', '请生成包含上述对话核心内容的摘要，保留重要事实和决策')
    ])
    summarization_chain = summarize_prompt | llm
    summary_message = summarization_chain.invoke({'chat_history': messages_to_summarize})

    chat_history.clear()
    chat_history.add(summary_message)
    for msg in last_two_messages:
        chat_history.add(msg)
    return True


result = chain_with_message_history.invoke({'input': '你好我是小白'},
                                           config=RunnableConfig(configurable={"session_id": "userTest"}))
print(result)
result2 = chain_with_message_history.invoke({'input': '我的名字叫什么'},
                                           config=RunnableConfig(configurable={"session_id": "userTest"}))
print(result2)

with gr.Blocks(title="姬", theme=gr.themes.Soft()) as block:
    pass

if __name__ == '__main__':
    block.launch()